import os
from PIL import Image
import numpy as np
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
import matplotlib.pyplot as plt
import torch
import random
from tqdm import tqdm
import argparse
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

from utils.common import check_and_create_directory, list_all_files, chunk_list
from utils.auto_sam import auto_sam
from utils.vit.vit import vit_encode

# 优化一下存储内容，emm 因为每张图单独存一个pt太大了
if __name__ == '__main__':
    
    path = '/mnt/g/tmp/'
    out_path = '/home/zry/datasets/building/train/tmp_bak'
    check_and_create_directory(out_path)

    fs = list_all_files(path, '.pt')
    for i, f in tqdm(enumerate(fs), total=len(fs)):
        if i< 115: continue
        masks = torch.load(f)
        new_pt = torch.ones(masks[0]['segmentation'].shape) * -1
        for idx, mask in enumerate(masks):
            seg_mask = mask['segmentation']
            new_pt[seg_mask] = idx
        name = os.path.split(f)[-1]
        new_f = os.path.join(out_path, name)
        torch.save(new_pt, new_f)

